diff --git a/arch/arm/mach-uniphier/dram/umc-ph1-pro4.c b/arch/arm/mach-uniphier/dram/umc-ph1-pro4.c
index 4d976e3..261f7cf 100644
--- a/arch/arm/mach-uniphier/dram/umc-ph1-pro4.c
+++ b/arch/arm/mach-uniphier/dram/umc-ph1-pro4.c
@@ -13,6 +13,16 @@
 #include "ddrphy-regs.h"
 #include "umc-regs.h"
 
+enum dram_size {
+	DRAM_SZ_128M,
+	DRAM_SZ_256M,
+	DRAM_SZ_512M,
+	DRAM_SZ_NR,
+};
+
+static u32 umc_initctlb[DRAM_SZ_NR] = {0x43030d3f, 0x43030d3f, 0x7b030d3f};
+static u32 umc_spcctla[DRAM_SZ_NR] = {0x002b0617, 0x003f0617, 0x00770617};
+
 static void umc_start_ssif(void __iomem *ssif_base)
 {
 	writel(0x00000001, ssif_base + 0x0000b004);
@@ -56,19 +66,36 @@
 	writel(0x00000001, ssif_base + UMC_DMDRST);
 }
 
-static void umc_dramcont_init(void __iomem *dramcont, void __iomem *ca_base,
-			      int size, int freq)
+static int umc_dramcont_init(void __iomem *dramcont, void __iomem *ca_base,
+			     int size, int width)
 {
+	enum dram_size dram_size;
+
+	switch (size / (width / 16)) {
+	case SZ_128M:
+		dram_size = DRAM_SZ_128M;
+		break;
+	case SZ_256M:
+		dram_size = DRAM_SZ_256M;
+		break;
+	case SZ_512M:
+		dram_size = DRAM_SZ_512M;
+		break;
+	default:
+		printf("unsupported DRAM size\n");
+		return -EINVAL;
+	}
+
 	writel(0x66bb0f17, dramcont + UMC_CMDCTLA);
 	writel(0x18c6aa44, dramcont + UMC_CMDCTLB);
 	writel(0x5101387f, dramcont + UMC_INITCTLA);
-	writel(0x43030d3f, dramcont + UMC_INITCTLB);
+	writel(umc_initctlb[dram_size], dramcont + UMC_INITCTLB);
 	writel(0x00ff00ff, dramcont + UMC_INITCTLC);
 	writel(0x00000d71, dramcont + UMC_DRMMR0);
 	writel(0x00000006, dramcont + UMC_DRMMR1);
 	writel(0x00000298, dramcont + UMC_DRMMR2);
 	writel(0x00000000, dramcont + UMC_DRMMR3);
-	writel(0x003f0617, dramcont + UMC_SPCCTLA);
+	writel(umc_spcctla[dram_size], dramcont + UMC_SPCCTLA);
 	writel(0x00ff0008, dramcont + UMC_SPCCTLB);
 	writel(0x000c00ae, dramcont + UMC_RDATACTL_D0);
 	writel(0x000c00ae, dramcont + UMC_RDATACTL_D1);
@@ -90,9 +117,11 @@
 	writel(0x200a0a00, dramcont + UMC_SPCSETB);
 	writel(0x00010000, dramcont + UMC_SPCSETD);
 	writel(0x80000020, dramcont + UMC_DFICUPDCTLA);
+
+	return 0;
 }
 
-static int umc_init_sub(int freq, int size_ch0, int size_ch1)
+int ph1_pro4_umc_init(const struct uniphier_board_data *bd)
 {
 	void __iomem *ssif_base = (void __iomem *)UMC_SSIF_BASE;
 	void __iomem *ca_base0 = (void __iomem *)UMC_CA_BASE(0);
@@ -103,6 +132,12 @@
 	void __iomem *phy0_1 = (void __iomem *)DDRPHY_BASE(0, 1);
 	void __iomem *phy1_0 = (void __iomem *)DDRPHY_BASE(1, 0);
 	void __iomem *phy1_1 = (void __iomem *)DDRPHY_BASE(1, 1);
+	int ret;
+
+	if (bd->dram_freq != 1600) {
+		pr_err("Unsupported DDR configuration\n");
+		return -EINVAL;
+	}
 
 	umc_dram_init_start(dramcont0);
 	umc_dram_init_start(dramcont1);
@@ -111,52 +146,43 @@
 
 	writel(0x00000101, dramcont0 + UMC_DIOCTLA);
 
-	ph1_pro4_ddrphy_init(phy0_0, freq, size_ch0);
+	ph1_pro4_ddrphy_init(phy0_0, bd->dram_freq, bd->dram_ch0_size);
 
 	ddrphy_prepare_training(phy0_0, 0);
 	ddrphy_training(phy0_0);
 
 	writel(0x00000103, dramcont0 + UMC_DIOCTLA);
 
-	ph1_pro4_ddrphy_init(phy0_1, freq, size_ch0);
+	ph1_pro4_ddrphy_init(phy0_1, bd->dram_freq, bd->dram_ch0_size);
 
 	ddrphy_prepare_training(phy0_1, 1);
 	ddrphy_training(phy0_1);
 
 	writel(0x00000101, dramcont1 + UMC_DIOCTLA);
 
-	ph1_pro4_ddrphy_init(phy1_0, freq, size_ch1);
+	ph1_pro4_ddrphy_init(phy1_0, bd->dram_freq, bd->dram_ch1_size);
 
 	ddrphy_prepare_training(phy1_0, 0);
 	ddrphy_training(phy1_0);
 
 	writel(0x00000103, dramcont1 + UMC_DIOCTLA);
 
-	ph1_pro4_ddrphy_init(phy1_1, freq, size_ch1);
+	ph1_pro4_ddrphy_init(phy1_1, bd->dram_freq, bd->dram_ch1_size);
 
 	ddrphy_prepare_training(phy1_1, 1);
 	ddrphy_training(phy1_1);
 
-	umc_dramcont_init(dramcont0, ca_base0, size_ch0, freq);
-	umc_dramcont_init(dramcont1, ca_base1, size_ch1, freq);
+	ret = umc_dramcont_init(dramcont0, ca_base0, bd->dram_ch0_size,
+				bd->dram_ch0_width);
+	if (ret)
+		return ret;
+
+	ret = umc_dramcont_init(dramcont1, ca_base1, bd->dram_ch1_size,
+				bd->dram_ch1_width);
+	if (ret)
+		return ret;
 
 	umc_start_ssif(ssif_base);
 
 	return 0;
 }
-
-int ph1_pro4_umc_init(const struct uniphier_board_data *bd)
-{
-	if (((bd->dram_ch0_size == SZ_512M && bd->dram_ch0_width == 32) ||
-	     (bd->dram_ch0_size == SZ_256M && bd->dram_ch0_width == 16)) &&
-	    ((bd->dram_ch1_size == SZ_512M && bd->dram_ch1_width == 32) ||
-	     (bd->dram_ch1_size == SZ_256M && bd->dram_ch1_width == 16)) &&
-	    bd->dram_freq == 1600) {
-		return umc_init_sub(bd->dram_freq,
-				    bd->dram_ch0_size / SZ_128M,
-				    bd->dram_ch1_size / SZ_128M);
-	} else {
-		pr_err("Unsupported DDR configuration\n");
-		return -EINVAL;
-	}
-}